{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# `check_basic_block_normal_form`\n",
    "\n",
    "参考：`tvm/tests/python/relay/test_analysis_basic_block_normal_form.py`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pytest\n",
    "import tvm\n",
    "import tvm.testing\n",
    "from tvm import relay\n",
    "from tvm.relay.analysis import check_basic_block_normal_form"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[31mERROR: usage: ipykernel_launcher.py [options] [file_or_dir] [file_or_dir] [...]\n",
      "ipykernel_launcher.py: error: unrecognized arguments: --f=/home/ai/.local/share/jupyter/runtime/kernel-v387e69e27796ecf682dbaa0af5c78919530db467c.json\n",
      "  inifile: /media/pc/data/lxw/ai/tvm-book/pyproject.toml\n",
      "  rootdir: /media/pc/data/lxw/ai/tvm-book\n",
      "\u001b[0m\n"
     ]
    },
    {
     "ename": "SystemExit",
     "evalue": "4",
     "output_type": "error",
     "traceback": [
      "An exception has occurred, use %tb to see the full traceback.\n",
      "\u001b[0;31mSystemExit\u001b[0m\u001b[0;31m:\u001b[0m 4\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/media/pc/data/lxw/envs/anaconda3a/envs/ai/lib/python3.12/site-packages/IPython/core/interactiveshell.py:3585: UserWarning: To exit: use 'exit', 'quit', or Ctrl-D.\n",
      "  warn(\"To exit: use 'exit', 'quit', or Ctrl-D.\", stacklevel=1)\n"
     ]
    }
   ],
   "source": [
    "\n",
    "def test_one_block():\n",
    "    x = relay.var(\"x\")\n",
    "    y = relay.add(x, x)\n",
    "    z = relay.add(x, y)\n",
    "    check_basic_block_normal_form(z)\n",
    "\n",
    "\n",
    "def test_let():\n",
    "    x = relay.var(\"x\")\n",
    "    y = relay.var(\"y\")\n",
    "    body = relay.Let(y, x, y)\n",
    "    check_basic_block_normal_form(body)\n",
    "\n",
    "\n",
    "@pytest.mark.xfail(raises=tvm.error.TVMError)\n",
    "def test_invalid_if():\n",
    "    cond = relay.var(\"cond\", dtype=\"bool\", shape=())\n",
    "    shared = relay.var(\"shared\")\n",
    "    true_branch = shared\n",
    "    false_branch = relay.add(shared, shared)\n",
    "    body = relay.If(cond, true_branch, false_branch)\n",
    "    \"\"\"\n",
    "    The program below violates basic block normal form, as the scope of %shared\n",
    "    is ambiguous and should not be in that of true branch.\n",
    "\n",
    "    free_var %cond: bool\n",
    "    if (%cond) {\n",
    "      free_var %shared\n",
    "      %shared\n",
    "    } else {\n",
    "      add(%shared, %shared)\n",
    "    }\n",
    "    \"\"\"\n",
    "    check_basic_block_normal_form(body)\n",
    "\n",
    "\n",
    "def test_valid_if():\n",
    "    cond = relay.var(\"cond\", dtype=\"bool\", shape=())\n",
    "    shared = relay.var(\"shared\")\n",
    "    true_branch = shared\n",
    "    false_branch = relay.add(shared, shared)\n",
    "    body = relay.If(cond, true_branch, false_branch)\n",
    "    shared_bound = relay.var(\"shared_bound\", shape=(1,), dtype=\"float32\")\n",
    "    body = relay.Let(shared, shared_bound, body)\n",
    "    \"\"\"\n",
    "    The program below uses let binding to control the scope of %shared, which\n",
    "    follows the basic block normal form.\n",
    "\n",
    "    free_var %shared_bound: Tensor[(1), float32]\n",
    "    let %shared = %shared_bound;\n",
    "    free_var %cond: bool\n",
    "    if (%cond) {\n",
    "      %shared\n",
    "    } else {\n",
    "      add(%shared, %shared)\n",
    "    }\n",
    "    \"\"\"\n",
    "    check_basic_block_normal_form(body)\n",
    "\n",
    "\n",
    "@pytest.mark.xfail(raises=tvm.error.TVMError)\n",
    "def test_invalid_if2():\n",
    "    \"\"\"\n",
    "    fn (%x: float32) {\n",
    "      %0 = equal(%x, 2f);\n",
    "      if (%0) {\n",
    "        %1 = add(%x, 1f);\n",
    "        multiply(%1, 2f)\n",
    "      } else {\n",
    "        multiply(%1, 1f)\n",
    "      }\n",
    "    }\n",
    "    \"\"\"\n",
    "    x = relay.var(\"x\", shape=(), dtype=\"float32\")\n",
    "    one = relay.const(1, dtype=\"float32\")\n",
    "    two = relay.const(2, dtype=\"float32\")\n",
    "    v1 = relay.add(x, one)\n",
    "    v2 = relay.equal(x, two)\n",
    "    true_branch = relay.multiply(v1, two)\n",
    "    false_branch = relay.multiply(v1, one)\n",
    "    body = relay.If(v2, true_branch, false_branch)\n",
    "    func = relay.Function([x], body)\n",
    "    check_basic_block_normal_form(func)\n",
    "\n",
    "\n",
    "def test_valid_if2():\n",
    "    \"\"\"\n",
    "    fn (%x: float32) {\n",
    "      let %v1 = add(%x, 1f);\n",
    "      %0 = equal(%x, 2f);\n",
    "      if (%0) {\n",
    "        multiply(%v1, 2f)\n",
    "      } else {\n",
    "        multiply(%v1, 1f)\n",
    "      }\n",
    "    }\n",
    "    \"\"\"\n",
    "    x = relay.var(\"x\", shape=(), dtype=\"float32\")\n",
    "    one = relay.const(1, dtype=\"float32\")\n",
    "    two = relay.const(2, dtype=\"float32\")\n",
    "    v1 = relay.var(\"v1\")\n",
    "    v2 = relay.equal(x, two)\n",
    "    true_branch = relay.multiply(v1, two)\n",
    "    false_branch = relay.multiply(v1, one)\n",
    "    body = relay.If(v2, true_branch, false_branch)\n",
    "    body = relay.Let(v1, relay.add(x, one), body)\n",
    "    func = relay.Function([x], body)\n",
    "    check_basic_block_normal_form(func)\n",
    "\n",
    "\n",
    "@pytest.mark.xfail(raises=tvm.error.TVMError)\n",
    "def test_func():\n",
    "    x = relay.var(\"x\", shape=(1,), dtype=\"float32\")  # , a)\n",
    "    y = relay.var(\"y\", shape=(1,), dtype=\"float32\")  # , a)\n",
    "    z = relay.var(\"z\", shape=(1,), dtype=\"float32\")  # , a)\n",
    "    x2 = relay.add(x, x)\n",
    "    func_a = relay.Function([y], relay.add(x2, y))  # , a, [a])\n",
    "    func_b = relay.Function([z], relay.add(x2, z))  # , a, [a])\n",
    "    body = relay.Tuple([func_a, func_b])\n",
    "    body = relay.Function([x], body)\n",
    "    \"\"\"\n",
    "    fn (%x: Tensor[(1), float32]) {\n",
    "      %1 = fn (%y: Tensor[(1), float32]) {\n",
    "        %0 = add(%x, %x);\n",
    "        add(%0, %y)\n",
    "      };\n",
    "      %2 = fn (%z: Tensor[(1), float32]) {\n",
    "        add(%0, %z)\n",
    "      };\n",
    "      (%1, %2)\n",
    "    }\n",
    "    \"\"\"\n",
    "    check_basic_block_normal_form(body)\n",
    "\n",
    "\n",
    "@pytest.mark.xfail(raises=tvm.error.TVMError)\n",
    "def test_higher_order_return():\n",
    "    x = relay.var(\"x\", shape=(1,), dtype=\"float32\")  # , a)\n",
    "    y = relay.var(\"y\", shape=(1,), dtype=\"float32\")  # , a)\n",
    "    z = relay.var(\"z\", shape=(1,), dtype=\"float32\")  # , a)\n",
    "    x2 = relay.add(x, x)\n",
    "    func_a = relay.Function([y], relay.add(x2, y))  # , a, [a])\n",
    "    func_b = relay.Function([z], relay.add(x2, z))  # , a, [a])\n",
    "    body = relay.Tuple([func_a, func_b])\n",
    "    body = relay.Function([x], body)\n",
    "    \"\"\"\n",
    "    fn (%x: Tensor[(1), float32]) {\n",
    "      %1 = fn (%y: Tensor[(1), float32]) {\n",
    "        %0 = add(%x, %x);\n",
    "        add(%0, %y)\n",
    "      };\n",
    "      %2 = fn (%z: Tensor[(1), float32]) {\n",
    "        add(%0, %z)\n",
    "      };\n",
    "      (%1, %2)\n",
    "    }\n",
    "    \"\"\"\n",
    "    check_basic_block_normal_form(body)\n",
    "\n",
    "\n",
    "@pytest.mark.xfail(raises=tvm.error.TVMError)\n",
    "def test_higher_order_nested():\n",
    "    x = relay.var(\"x\", dtype=\"float32\", shape=(1,))\n",
    "    s = relay.var(\"s\", dtype=\"float32\", shape=(1,))\n",
    "    shared = relay.add(s, s)\n",
    "    func_true = relay.Function([x], relay.add(x, shared))\n",
    "    choice_t = relay.FuncType([], relay.scalar_type(\"bool\"))\n",
    "    f = relay.Var(\"f\", choice_t)\n",
    "    z = relay.Var(\"z\")\n",
    "    body = relay.If(f(), func_true, relay.Function([z], relay.add(z, shared)))\n",
    "    top = relay.Function([f, s], body)\n",
    "    \"\"\"\n",
    "    fn (%f: fn () -> bool, %s: Tensor[(1), float32]) {\n",
    "      %0 = %f();\n",
    "      if (%0) {\n",
    "        fn (%x: Tensor[(1), float32]) {\n",
    "          %1 = add(%s, %s);\n",
    "          add(%x, %1)\n",
    "        }\n",
    "      } else {\n",
    "        fn (%z) {\n",
    "          add(%z, %1)\n",
    "        }\n",
    "      }\n",
    "    }\n",
    "    \"\"\"\n",
    "    check_basic_block_normal_form(top)\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    tvm.testing.main()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ai",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
